#include "../../include/core/container.h"
#include <cstring>
#include <iostream>

#define INPUT_SIZE 16
#define OUTPUT_SIZE 16
#define KEY_SIZE 16

extern "C" {
void AES_CRACK_OP(MemRef<unsigned char, 1> *inputs, int msgLen,
                  MemRef<unsigned char, 1> *outputs, int cptLen,
                  MemRef<unsigned char, 1> *key, int keyLen, int msgCnt);
}

uint8_t exampleInput[INPUT_SIZE] = {0xd0, 0x41, 0x30, 0xd3, 0xee, 0xd2,
                                    0xbc, 0xe7, 0xc2, 0x21, 0x87, 0xc8,
                                    0x5b, 0x09, 0xc7, 0x75};
uint8_t exampleKey[KEY_SIZE] = {'A', '0', 'B', '1', 'C', '2', 'D', '3',
                                'E', '4', 'F', '5', 'G', '6', 'H', '7'};

MemRef<unsigned char, 1> input{exampleInput, INPUT_SIZE};
MemRef<unsigned char, 1> key{exampleKey, KEY_SIZE};
uint8_t *exampleOutput = new uint8_t[OUTPUT_SIZE];
MemRef<unsigned char, 1> output{exampleOutput, OUTPUT_SIZE};

void printResult(uint8_t *output) {
  // Print the output.
  std::cout << "[ ";
  for (size_t i = 0; i < OUTPUT_SIZE; i += 4) {
    for(int j = 3; j >=0; j--){
      printf("%c",*(output+i+j));
    }
  }
  std::cout << " ]" << std::endl;
}

void printInput(uint8_t *input) {
  // Print the input.
  std::cout << "AES Input : [ ";
  for (size_t i = 0; i < INPUT_SIZE; i++) {
    printf("%02x", input[i]);
  }
  std::cout << " ] ==> ";
}

int main() {
  int keyLen = 16;
  AES_CRACK_OP(&input, INPUT_SIZE, &output, 0, &key, 0, 1);
  printInput(exampleInput);
  printResult(exampleOutput);

  std::cout << "Key " << ": [";
  for (size_t i = 0; i < keyLen; i++) {
    std::cout << key[i];
  }
  std::cout << "]" << std::endl;

  return 0;
}
